{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o4PAQZyTegH3"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/cohere-ai/notebooks/blob/main/notebooks/llmu/Introduction_to_RAG.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EimVurhQ45yk"
      },
      "source": [
        "# Introduction to RAG"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4jBEKety4mKZ"
      },
      "source": [
        "This notebook shows a quickstart example on how to build a RAG-powered chatbot with the Cohere's Chat endpoint. The chatbot can extract relevant information from external documents and produce verifiable, inline citations in its responses.\n",
        "\n",
        "The diagram below provides an overview of what we’ll build.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2rzGO7hVtttt"
      },
      "source": [
        "![rag-workflow-1.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Qci1A1RNWSU"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KVhenxwXE9oI",
        "outputId": "e1363e93-86cc-407a-d347-0a88d399e955"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/118.8 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m118.8/118.8 kB\u001b[0m \u001b[31m3.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.6/75.6 kB\u001b[0m \u001b[31m8.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.8/77.8 kB\u001b[0m \u001b[31m6.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m6.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "! pip install cohere -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "KiAF_FuHF2Mg",
        "outputId": "2c907cf4-bea5-484a-8e7d-dfe1948a4922"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "#@title Enable text wrapping in Google Colab\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2gY7O7rqE6Nq"
      },
      "outputs": [],
      "source": [
        "import cohere\n",
        "co = cohere.Client(\"COHERE_API_KEY\") # Get your API key here: https://dashboard.cohere.com/api-keys"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JLXxYR_UHGLJ"
      },
      "source": [
        "# Define documents"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bKbrq5c1e45y"
      },
      "source": [
        "We define the documents that we want to ground an LLM’s response with, formatted as a list. In our case, each document consists of two fields: title and text.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CdxeI3XW4yIH"
      },
      "outputs": [],
      "source": [
        "documents = [\n",
        "    {\n",
        "        \"title\": \"Tall penguins\",\n",
        "        \"text\": \"Emperor penguins are the tallest.\"},\n",
        "    {\n",
        "        \"title\": \"Penguin habitats\",\n",
        "        \"text\": \"Emperor penguins only live in Antarctica.\"},\n",
        "    {\n",
        "        \"title\": \"What are animals?\",\n",
        "        \"text\": \"Animals are different from plants.\"}\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bT3KjYATHH9o"
      },
      "source": [
        "# Generate response with citations"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A64u2fIyfBri"
      },
      "source": [
        "Cohere’s RAG functionalities are part of the Chat endpoint, with the Command model as the underlying LLM. This allows developers to build chatbots that have the full context of a conversation and are not limited to a single interaction.\n",
        "\n",
        "First, we define the user message. Then we generate the response from the LLM and display it, together with citations and the source documents used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fkZ6gEYGISWZ",
        "outputId": "662fe3ee-6b61-487b-e59c-c6a18e6038ce"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The tallest living species of penguin is the emperor penguin (Apteryx australis), which can measure up to 1.6 to 1.8 m (5.2 to 6 ft) when fully grown.\n",
            "\n",
            "CITATIONS:\n",
            "start=45 end=60 text='emperor penguin' document_ids=['doc_0']\n",
            "start=61 end=80 text='(Apteryx australis)' document_ids=['doc_0']\n",
            "start=106 end=118 text='1.6 to 1.8 m' document_ids=['doc_0']\n",
            "start=119 end=132 text='(5.2 to 6 ft)' document_ids=['doc_0']\n",
            "\n",
            "DOCUMENTS:\n",
            "{'id': 'doc_0', 'text': 'Emperor penguins are the tallest.', 'title': 'Tall penguins'}\n",
            "{'id': 'doc_1', 'text': 'Emperor penguins only live in Antarctica.', 'title': 'Penguin habitats'}\n",
            "{'id': 'doc_2', 'text': 'Animals are different from plants.', 'title': 'What are animals?'}\n"
          ]
        }
      ],
      "source": [
        "# Get the user message\n",
        "message = \"What are the tallest living penguins?\"\n",
        "\n",
        "# Generate the response\n",
        "response = co.chat_stream(message=message,\n",
        "                          documents=documents)\n",
        "\n",
        "\n",
        "# Display the response\n",
        "citations = []\n",
        "cited_documents = []\n",
        "\n",
        "for event in response:\n",
        "    if event.event_type == \"text-generation\":\n",
        "        print(event.text, end=\"\")\n",
        "    elif event.event_type == \"citation-generation\":\n",
        "        citations.extend(event.citations)\n",
        "    elif event.event_type == \"search-results\":\n",
        "        cited_documents = event.documents\n",
        "\n",
        "# Display the citations and source documents\n",
        "if citations:\n",
        "  print(\"\\n\\nCITATIONS:\")\n",
        "  for citation in citations:\n",
        "    print(citation)\n",
        "\n",
        "  print(\"\\nDOCUMENTS:\")\n",
        "  for document in cited_documents:\n",
        "    print(document)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
